#导入相应的库
from tensorflow.keras.preprocessing.image import ImageDataGenerator
import tensorflow as tf
import os
from tensorflow import keras
from tensorflow.compat.v1 import ConfigProto
from tensorflow.compat.v1 import InteractiveSession

config = ConfigProto()
config.gpu_options.allow_growth = True
session = InteractiveSession(config=config)


os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'


#设置图片的高和宽，一次训练所选取的样本数，迭代次数
im_height = 224
im_width = 224
batch_size = 64
epochs = 500

train_dir = './train_data' # 训练集路径

# 定义训练集图像生成器，并进行图像增强
train_image_generator = ImageDataGenerator( rescale=1./255,  # 归一化
                                            rotation_range=40, #旋转范围
                                            width_shift_range=0.2, #水平平移范围
                                            height_shift_range=0.2, #垂直平移范围
                                            shear_range=0.2, #剪切变换的程度
                                            zoom_range=0.2, #缩放范围
                                            horizontal_flip=True,  #水平翻转
                                            fill_mode='nearest')  
                                             
# 使用图像生成器从文件夹train_dir中读取样本，对标签进行one-hot编码
train_data_gen = train_image_generator.flow_from_directory(directory=train_dir,
                                                           batch_size=batch_size,
                                                           shuffle=True,   #打乱数据
                                                           target_size=(im_height, im_width),
                                                           class_mode='categorical')
# 训练集样本数    
# total_train = train_data_gen.n
# print(train_data_gen.class_indices)


logdir = './callbacks'  # window里用os.path.join('callbacks')
if not os.path.exists(logdir):
    os.mkdir(logdir)
output_model_file = os.path.join(logdir,
                                 "MobilenetV2_model.h5")

# 要写分布式 下次记得写成 main函数的形式  多进程要在main里面调用
# 分布式不支持fit_generator  所以要用fit  fit会一次性把数据加载到内存
# strategy = tf.distribute.MirroredStrategy()  


covn_base = tf.keras.applications.MobileNetV2(weights='imagenet',include_top=False)

#构建模型    
model = tf.keras.Sequential()
model.add(covn_base)
model.add(tf.keras.layers.GlobalAveragePooling2D()) #加入全局平均池化层
model.add(tf.keras.layers.Dense(725,activation='softmax')) #加入输出层(725分类)

#编译模型
model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.001),  #使用adam优化器，学习率为0.0001
            loss=tf.keras.losses.CategoricalCrossentropy(from_logits=False), #交叉熵损失函数
            metrics=["accuracy"])  #评价函数

model.summary()

callbacks = [
    # keras.callbacks.TensorBoard(logdir),
    keras.callbacks.ModelCheckpoint(output_model_file,
                                    save_best_only = True),
    # keras.callbacks.EarlyStopping(monitor='val_acc',patience=5, min_delta=1e-3),
]

#开始训练
model.fit_generator(train_data_gen,validation_data=train_data_gen,epochs=epochs,callbacks=callbacks,workers=32)
